import torch
import torch.nn as nn
import torch.nn.functional as F

class ResidualBlock(nn.Module):
    """残差块"""
    def __init__(self, in_channels, out_channels, stride=1):
        super(ResidualBlock, self).__init__()
        
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm1d(out_channels)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm1d(out_channels)
            )
            
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class HARResNet(nn.Module):
    """基于ResNet的人类活动识别模型"""
    def __init__(self, input_size=561, num_classes=6, dropout=0.5):
        super(HARResNet, self).__init__()
        
        # 将输入特征重塑为2D形式以用于CNN
        self.channels = 3  # 假设为三轴加速度计数据
        self.seq_length = input_size // self.channels
        
        # 初始卷积层
        self.conv1 = nn.Conv1d(self.channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm1d(64)
        self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
        
        # 残差层
        self.layer1 = self._make_layer(64, 64, 2, stride=1)
        self.layer2 = self._make_layer(64, 128, 2, stride=2)
        self.layer3 = self._make_layer(128, 256, 2, stride=2)
        
        # 分类器 - 使用自适应池化，避免固定输出大小的问题
        self.fc = nn.Sequential(
            nn.Linear(256, 512),  # 只使用通道数作为输入特征维度
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(512, num_classes)
        )
        
    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
        layers = []
        layers.append(ResidualBlock(in_channels, out_channels, stride))
        for _ in range(1, num_blocks):
            layers.append(ResidualBlock(out_channels, out_channels))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        batch_size = x.size(0)
        
        # 重塑输入以适应CNN
        x = x.view(batch_size, self.channels, -1)
        
        # 特征提取
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        
        # 全局平均池化 - 将每个通道池化为单个值
        x = F.adaptive_avg_pool1d(x, 1)
        x = x.view(batch_size, -1)  # 展平为 [batch_size, 256]
        
        # 分类
        x = self.fc(x)
        
        return x